# flake8: noqa: E501
import copy
import io
from contextlib import redirect_stdout
from typing import Any, Optional, Type

from asyncer import asyncify

from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
from lagent.actions.parser import BaseParser, JsonParser
from lagent.schema import ActionReturn, ActionStatusCode


class GenericRuntime:
    """通用 Python 运行时环境"""
    
    GLOBAL_DICT = {}  # 全局变量字典
    LOCAL_DICT = None  # 局部变量字典
    HEADERS = []  # 预加载代码列表

    def __init__(self):
        # 复制全局和局部变量环境
        self._global_vars = copy.copy(self.GLOBAL_DICT)
        self._local_vars = copy.copy(self.LOCAL_DICT) if self.LOCAL_DICT else None

        # 执行预加载代码
        for c in self.HEADERS:
            self.exec_code(c)

    def exec_code(self, code_piece: str) -> None:
        """执行代码片段"""
        exec(code_piece, self._global_vars)

    def eval_code(self, expr: str) -> Any:
        """计算表达式的值"""
        return eval(expr, self._global_vars)


class PythonInterpreter(BaseAction):
    """
    Python 代码执行器。
    
    Args:
        answer_symbol: LLM 的答案符号
        answer_expr: Python 脚本的答案函数名，默认为 'solution()'
        answer_from_stdout: 是否从标准输出获取执行结果
        timeout: 执行超时时间(秒)，默认 20 秒
        description: 工具描述
        parser: 结果解析器类
    """

    def __init__(
        self,
        answer_symbol: Optional[str] = None,
        answer_expr: Optional[str] = 'solution()',
        answer_from_stdout: bool = False,
        timeout: int = 20,
        description: Optional[dict] = None,
        parser: Type[BaseParser] = JsonParser,
    ) -> None:
        super().__init__(description, parser)
        self.answer_symbol = answer_symbol
        self.answer_expr = answer_expr
        self.answer_from_stdout = answer_from_stdout
        self.timeout = timeout

    @tool_api
    def run(self, command: str) -> ActionReturn:
        """
        执行 Python 代码。代码必须定义一个名为 'solution' 的函数。

        代码格式示例:
        ```python
        # 导入依赖包
        import xxx
        
        def solution():
            # 初始化变量
            variable = xxx
            # 步骤一
            mid_result = func(variable)
            # 步骤 x
            mid_result = func(mid_result)
            # 最终结果
            final_answer = func(mid_result)
            return final_answer
        ```

        Args:
            command: Python 代码片段
        """
        from func_timeout import FunctionTimedOut, func_set_timeout

        # 创建运行时环境
        self.runtime = GenericRuntime()
        try:
            # 带超时执行代码
            tool_return = func_set_timeout(self.timeout)(self._call)(command)
        except FunctionTimedOut as e:
            # 处理超时错误
            tool_return = ActionReturn(type=self.name)
            tool_return.errmsg = repr(e)
            tool_return.state = ActionStatusCode.API_ERROR
        return tool_return

    def _call(self, command: str) -> ActionReturn:
        """执行代码并获取结果"""
        tool_return = ActionReturn(type=self.name)
        try:
            # 提取代码块
            if '```python' in command:
                command = command.split('```python')[1].split('```')[0]
            elif '```' in command:
                command = command.split('```')[1].split('```')[0]
            tool_return.args = dict(text='```python\n' + command + '\n```')
            command = command.split('\n')

            # 根据不同模式获取结果
            if self.answer_from_stdout:
                # 从标准输出获取结果
                program_io = io.StringIO()
                with redirect_stdout(program_io):
                    self.runtime.exec_code('\n'.join(command))
                program_io.seek(0)
                res = program_io.readlines()[-1]
            elif self.answer_symbol:
                # 从全局变量获取结果
                self.runtime.exec_code('\n'.join(command))
                res = self.runtime._global_vars[self.answer_symbol]
            elif self.answer_expr:
                # 从表达式获取结果
                self.runtime.exec_code('\n'.join(command))
                res = self.runtime.eval_code(self.answer_expr)
            else:
                # 执行代码并计算最后一行
                self.runtime.exec_code('\n'.join(command[:-1]))
                res = self.runtime.eval_code(command[-1])
        except Exception as e:
            # 处理执行错误
            tool_return.errmsg = repr(e)
            tool_return.type = self.name
            tool_return.state = ActionStatusCode.API_ERROR
            return tool_return

        try:
            # 格式化结果
            tool_return.result = [dict(type='text', content=str(res))]
            tool_return.state = ActionStatusCode.SUCCESS
        except Exception as e:
            # 处理格式化错误
            tool_return.errmsg = repr(e)
            tool_return.type = self.name
            tool_return.state = ActionStatusCode.API_ERROR
        return tool_return


class AsyncPythonInterpreter(AsyncActionMixin, PythonInterpreter):
    """异步版本的 Python 代码执行器"""
    
    @tool_api
    @asyncify  # 将同步方法转换为异步方法
    def run(self, command: str) -> ActionReturn:
        """异步执行 Python 代码"""
        return super().run(command)
